#!/bin/bash
# File Name   : cg_manager.sh
# Author      : moshan
# Mail        : mo_shan@yeah.net
# Version     : 1.0
# Created Time: 2025-03-12 11:08:00
# Function    : cgroup规则管理
#########################################################################
set -e -o pipefail

work_dir="/opt/soft/cgroup"
conf_dir="${work_dir}/conf"

localhost_ip="$(/sbin/ip a 2>/dev/null|grep inet|grep brd|grep scope|awk '{print $2}'|awk -F'/' '{print $1}'|head -1)" || true

. /etc/profile || true

env="${1}"

export PATH="/opt/soft/mysql80/bin:/opt/soft/mysql57/bin:${PATH}"

declare -A type_file_map 

type_file_map=(
    ["cpu"]="cpu.cfs_quota_us"
    ["mem"]="memory.limit_in_bytes"
    ["io_read"]="blkio.throttle.read_bps_device"
    ["io_write"]="blkio.throttle.write_bps_device"
    ["iops_read"]="blkio.throttle.read_iops_device"
    ["iops_write"]="blkio.throttle.write_iops_device"
)

function f_get_disk_num()
{
    declare mp

    mp="${1}"

    lsblk -o NAME,MAJ:MIN,RM,SIZE,RO,TYPE,MOUNTPOINT | awk -v mp="${mp}" -F "[[:space:]]+" '
    
    NR == 1 { next; }
    
    $6 == "disk" {
    
        disk = $1;
    
        num[disk] = $2;
    
        if ($NF) {
    
            gsub(/^[[:space:]]+|[[:space:]]+$/, "", $NF);
    
            if($NF==mp) {print(num[disk])};
    
        }
    }
    
    $6 == "part" && $NF {
    
        gsub(/^[[:space:]]+|[[:space:]]+$/, "", $NF);
    
        if($NF==mp) {print(num[disk])};
    
    }'  
}

function f_cpu_manager()
{
    declare group mvalue pid opt
    group="${1}"
    mvalue="${2}"
    pid="${3}"
    opt="${4}"

    if [ "${opt}x" == "delx" ]
    then
        cgdelete -g cpu:${group}
    else
        cgcreate -g cpu:/${group}
        echo 100000 > /sys/fs/cgroup/cpu/${group}/cpu.cfs_period_us
        echo ${mvalue} > /sys/fs/cgroup/cpu/${group}/cpu.cfs_quota_us
        cgclassify -g cpu:${group} ${pid}
    fi
    return 0
}

function f_memory_manager()
{
    declare group mvalue mtype pid opt
    group="${1}"
    mvalue="${2}"
    pid="${3}"
    mtype="${4}"
    opt="${5}"

    if [ "${opt}x" == "delx" ]
    then
        cgdelete -g memory:${group}
    else
        cgcreate -g memory:/${group}
        echo "${mvalue}G" > /sys/fs/cgroup/memory/${group}/memory.limit_in_bytes
        cgclassify -g memory:${group} ${pid}

        if [ "${mtype}x" == "memx" ] || [ "${mtype}x" == "memoryx" ]
        then
            echo 1 > /sys/fs/cgroup/memory/${group}/memory.oom_control
        else
            echo 0 > /sys/fs/cgroup/memory/${group}/memory.oom_control
        fi
    fi

    return 0
}

function f_blkio_manager()
{
    declare group mvalue mtype pid opt
    group="${1}"
    mvalue="${2}"
    pid="${3}"
    mtype="${4}"
    diskid="${5}"
    opt="${6}"

    if [ "${opt}x" == "delx" ]
    then
        cgdelete -g blkio:${group}
    else

        [ -z "${diskid}" ] && return 11
        cgcreate -g blkio:/${group}
        echo "${diskid} ${mvalue}" >> /sys/fs/cgroup/blkio/${group}/${type_file_map["$4"]}
        cgclassify -g blkio:${group} ${pid}
    fi

    return 0
}

function f_net_manager()
{
    nic_name=$(/sbin/ip a|grep ${localhost_ip}|grep "scope global"|awk '{print $NF}')
    if [ "${4}x" == "net_outx" ]
    then

echo "
        cgcreate -g net_cls:/${1}
        echo 0x10000${5} > /sys/fs/cgroup/net_cls/${1}/net_cls.classid
        tc qdisc add dev ${nic_name} root handle 10: htb 2>/dev/null || true #每个机器只能有一个root规则
        tc class add dev ${nic_name} parent 10: classid 10:${5} htb rate ${2}mbit #这里可以带上ceil 300mbit 表示上限也是这个，当不带的时候，如果机器带宽有富裕可以不受限制
        tc filter add dev ${nic_name} protocol ip parent 10: prio 10 handle 0x10000${5} fw flowid 10:${5}
        cgclassify -g net_cls:${1} ${3}

        tc filter del dev ${nic_name} protocol ip parent 10: prio 10 handle 0x10000${5}
        tc class del dev ${nic_name} parent 10: classid 10:${5} 2>/dev/null
        tc qdisc del dev ${nic_name} root handle 10: htb 2>/dev/null
"
    else
echo "
        cgcreate -g net_cls:/${1}
        echo 0x10000${5} > /sys/fs/cgroup/net_cls/myapp/net_cls.classid
        modprobe ifb numifbs=$((${5}+1))
        ip link set dev ifb${5} up
        tc qdisc add dev ${nic_name} handle ffff: ingress #这是一个特殊的句柄值，专门用于标识 ingress 队列调度器
        tc filter add dev ${nic_name} parent ffff: protocol ip u32 match ip dport ${6} 0xffff action mirred egress redirect dev ifb${5} #掩码 0xffff 是一个十六进制值，转换成二进制是 1111111111111111，即所有位都为 1。这意味着：完全匹配
        tc qdisc add dev ifb${5} root handle 10: htb default 1
        tc class add dev ifb${5} parent 10: classid 10:${5} htb rate ${2}mbit ceil ${2}mbit
        tc filter add dev ifb${5} protocol ip parent 10: prio 1 u32 match u32 0x10000${5} 0xffffff at 16 flowid 10:${5}

        tc filter del dev ifb${5} protocol ip parent 10: prio 1 u32 match u32 0x10000${5} 0xffffff at 16 flowid 10:${5}
        tc class del dev ifb${5} parent 10: classid 10:${5} 2>/dev/null
        tc qdisc del dev ifb${5} root handle 10: htb 2>/dev/null
        tc filter del dev ${nic_name} parent ffff: protocol ip prio 1 u32 match ip dport ${6} 0xffff
        tc qdisc del dev ${nic_name} root handle 10: htb 2>/dev/null
        ip link set dev ifb${5} down
        ip link delete ifb${5}
        rmmod ifb

"
    fi
    return 0
}

function f_check_master()
{
    declare exe_state is_master
    exe_state="0"
    is_master="$(timeout 5s ${mysql_comm} -NBe "show slave status\G" 2>/dev/null|wc -l)" || exe_state="${?}"

    if [ "${exe_state}x" != "0x" ] 
    then
        exe_state="0"
        is_master="$(timeout 5s ${local_mysql_comm} -NBe "show slave status\G" 2>/dev/null|wc -l)" || exe_state="${?}"
        mysql_comm="${local_mysql_comm}"
    fi

    if [ "${is_master}x" != "0x" ] || [ "${exe_state}x" != "0x" ]
    then
        return 1
    fi
}

function f_main()
{
    [ -f "/opt/soft/cgroup/stop_file_long" ] && return 0

    if [ "${env}x" != "devx" ] && [ "${env}x" != "oltpx" ]
    then
        return 1
    else
        source ${conf_dir}/${env}.conf
    fi

    mysql_comm="${mysql_path} -u${mysql_user} -p${mysql_pass} -h${mysql_host} -P ${mysql_port} ${mysql_db}"
    local_mysql_comm="${mysql_path} -u${mysql_user} -p${mysql_pass} -h${localhost_ip} -P ${mysql_port} ${mysql_db}"

    f_check_master

    myker="$(uname -r|awk -F. '{print $1"."$2}')"

    ker_mark="false"
    for ker in $(echo "${kernel}"|tr "," " ")
    do
        if [ "${ker}x" == "${myker}x" ]
        then
            ker_mark="true"
            break
        fi
    done

    if [ "${ker_mark}x" == "falsex" ]
    then
        usql="update tb_cgroup_manager set istate = -1 , remarks = '${myker}不在${kernel}列表内,内核版本不支持' where rshost = '${localhost_ip}' and istate = 3"
        echo "${usql}"|timeout 60s ${mysql_comm} -NB 2>/dev/null
        return 1
    fi

    declare -a info_list

    sql="SELECT concat(rsport,'@@@',isotype,'@@@',isovalue,'@@@',istate,'@@@',pidpath,'@@@',mountpoin,'@@@',dbtype,'@@@',id,'@@@',rule) 
         from tb_cgroup_manager 
         where a_time > (DATE_SUB(NOW(), INTERVAL 60000 MINUTE)) and rshost = '${localhost_ip}' and istate = 3 order by rule desc;"

    info_list=($(echo "${sql}" |timeout 60s ${mysql_comm} -NB 2>/dev/null))

    for info in ${info_list[@]}
    do
        rsport="$(awk -F'@@@' '{print $1}' <<< "${info}")"
        isotype="$(awk -F'@@@' '{print $2}' <<< "${info}")"
        isovault="$(awk -F'@@@' '{print $3}' <<< "${info}")"
        istate="$(awk -F'@@@' '{print $4}' <<< "${info}")"
        pidpath="$(awk -F'@@@' '{print $5}' <<< "${info}")"

        if [ -f "${pidpath}" ]
        then
            pid="$(cat ${pidpath})"
        else
            which lsof0 >/dev/null 2>&1 || yum -y install lsof > /dev/null 2>&1
            pid="$(lsof -i:3574 -t)"
        fi

        which cgcreate >/dev/null 2>&1 || { chattr -i /etc/passwd /etc/group /etc/gshadow /etc/shadow && yum -y install libcgroup libcgroup-tools && chattr +i /etc/passwd /etc/group /etc/gshadow /etc/shadow;}

        mountpoin="$(awk -F'@@@' '{print $6}' <<< "${info}")"
        dbtype="$(awk -F'@@@' '{print $7}' <<< "${info}")"
        id="$(awk -F'@@@' '{print $8}' <<< "${info}")"
        rule="$(awk -F'@@@' '{print $9}' <<< "${info}")"

        cgroup="${dbtype}${rsport}"

        usql="update tb_cgroup_manager set istate = 2 where id = ${id} and istate = 3"
        usql_err="update tb_cgroup_manager set istate = -1 where id = ${id} and istate = 2"

        echo "${usql}"|timeout 60s ${mysql_comm} -NB 2>/dev/null

        exe_state="0"
        del_opt=""

        if [ "${isotype}x" == "mem_oomx" ] || [ "${isotype}x" == "memx" ] || [ "${isotype}x" == "memoryx" ] || [ "${isotype}x" == "memory_oomx" ]
        then
            del_opt="memory"
            [ ${isovault} -le 1 ] && isovault=1
            f_memory_manager "${cgroup}" "${isovault}" "${pid}" "${isotype}" "${rule}" || exe_state="${?}"
        elif [ "${isotype}x" == "cpux" ]
        then
            del_opt="cpu"
            [ ${isovault} -le 1 ] && isovault=1
            f_cpu_manager "${cgroup}" "$((${isovault}*100000))" "${pid}" "${rule}" || exe_state="${?}"
        elif [ "${isotype}x" == "io_readx" ] || [ "${isotype}x" == "io_writex" ]
        then
            del_opt="blkio"
            [ ${isovault} -le 1 ] && isovault=1
            f_blkio_manager "${cgroup}" "$((${isovault}*1024*1024))" "${pid}" "${isotype}" "$(f_get_disk_num "${mountpoin}")" "${rule}" || exe_state="${?}"
            [ ${exe_state} -eq 11 ] && usql_err="update tb_cgroup_manager set istate = -1,remarks = '数据盘挂载点不存在'  where id = ${id} and istate = 2"
        elif [ "${isotype}x" == "iops_readx" ] || [ "${isotype}x" == "iops_writex" ]
        then
            del_opt="blkio"
            [ ${isovault} -le 10 ] && isovault=10
            f_blkio_manager "${cgroup}" "${isovault}" "${pid}" "${isotype}" "$(f_get_disk_num "${mountpoin}")" "${rule}" || exe_state="${?}"
            [ ${exe_state} -eq 11 ] && usql_err="update tb_cgroup_manager set istate = -1,remarks = '数据盘挂载点不存在'  where id = ${id} and istate = 2"
        elif [ "${isotype}x" == "net_inx" ] || [ "${isotype}x" == "net_outx" ]
        then
            del_opt="net_cls"
            [ ${isovault} -le 100 ] && isovault=100
            f_net_manager "${cgroup}" "$((${isovault}*8))" "${pid}" "${isotype}" "$((${isovault}/100))" "${rsport}" || exe_state="${?}"
        else
            echo "${usql-err}"|timeout 60s ${mysql_comm} -NB 2>/dev/null
            continue
        fi

        if [ "${exe_state}x" != "0x" ]
        then
            cgdelete -g ${del_opt}:${group}
            echo "${usql-err}"|timeout 60s ${mysql_comm} -NB 2>/dev/null
            continue
        fi

        usql="update tb_cgroup_manager set istate = 1 where id = ${id} and istate = 2"
        if [ "${rule}x" == "delx" ]
        then
            usql="update tb_cgroup_manager set istate = 0 where id = ${id} and istate = 2"
        fi

        echo "${usql}"|timeout 60s ${mysql_comm} -NB 2>/dev/null
    done
}

f_main "${@}"
